from pykalman import KalmanFilter
import numpy as np

def Kalman1D(observations, damping=1.0):
    # To return the smoothed time series data
    observation_covariance = damping
    initial_value_guess = observations[0]
    transition_matrix = 1
    transition_covariance = 0.1
    kf = KalmanFilter(
        initial_state_mean=initial_value_guess,
        initial_state_covariance=observation_covariance,
        observation_covariance=observation_covariance,
        transition_covariance=transition_covariance,
        transition_matrices=transition_matrix
    )
    pred_state, state_cov = kf.smooth(observations)
    return pred_state


def missing_value_fix(l,fix_type = 1):
    # 计算判断异常点和极端异常点的临界值
    np_ar = np.array(l)
    outlier_ll = np_ar.mean() - 2 * np_ar.std()
    outlier_ul = np_ar.mean() + 2 * np_ar.std()

    extreme_outlier_ll = np_ar.mean() - 3 * np_ar.std()
    extreme_outlier_ul = np_ar.mean() + 3 * np_ar.std()
    errs = []
    if fix_type == 1:
        errs.extend((np.where(np_ar < extreme_outlier_ll))[0].tolist())
        errs.extend(np.where(np_ar > extreme_outlier_ul)[0].tolist())
    else:
        errs.extend((np.where(np_ar < outlier_ll))[0].tolist())
        errs.extend(np.where(np_ar > outlier_ul)[0].tolist())
    errs.sort()
    errs.reverse()
    for index in errs:
        repeat = index - 1
        while (repeat > 0):
            if repeat in errs:
                repeat -= 1
                continue
            l[index] = l[repeat]
            break
        if repeat <= 0:
            l[index] = l[index + 1]
        del index
    return l

def list_smooth(pls, is_kalman=True,is_missing_fix = True,kalman_d = 1,missing_type = 1):
    x_list = []
    y_list = []
    z_list = []
    new_pls = []
    for x, y, z in pls:
        x_list.append(x)
        y_list.append(y)
        z_list.append(z)
    if is_missing_fix:
        print('正在修复缺失值.....')
        x_list = missing_value_fix(x_list,missing_type)
        y_list = missing_value_fix(y_list,missing_type)
    if is_kalman :
        print('正在进行卡尔曼滤波.....')
        x_list = [i[0] for i in Kalman1D(x_list, kalman_d)]
        y_list = [i[0] for i in Kalman1D(y_list, kalman_d)]
    z_list_smooth = z_list
    for i in range(len(x_list)):
        new_pls.append((x_list[i], y_list[i], z_list_smooth[i]))
    return new_pls